from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

class imageCaptioning:
    # AI模型的本地地址
    modelName = 'damo/mplug_image-captioning_coco_base_zh'
    # 使用 cpu 还是 GPU 计算，GPU 比 CPU快，但是 GPU 只支持 英伟达 显卡，至少4G显存以上
    deviceType = 'cpu' # cpu, gpu
    #  外部传入 需要检索的关键词数组
    keyword = []
    # 外部传入 的视频帧 截图数据，[ 'http地址 或者 本地地址', '' ]
    # 地址可以用 video标签 poster 预览图，或者使用 ffmpeg，docarray，这里推荐 使用 多种模态数据结构工具包 docarray
    # 具体 docarray 封装看 utils/screenshot.py
    imgArr = []
    # 符合 条件 数据的计数
    count = 0
    # 符合 条件 数据的list下标
    haveIndex = []

    def __init__(self, key, imgList):
        self.keyword = key
        self.imgArr = imgList

    def start(self):
        self.count = 0
        self.haveIndex = []
    
        for index, item in enumerate(self.imgArr):
            print()
            print("ai 得到的图片为：", item)
            res = self.mplugImage(item)
            print("ai 解析图片内容为：", res)
            for v in self.keyword:
                if v in res:
                    self.count += 1
                    self.haveIndex.append(index)


        if self.count != 0 and len(self.haveIndex) != 0:
            # print('包含次数:', self.count)
            return {'count': self.count, 'index': self.haveIndex}
        else:
            return {'count': 0, 'index': []}

      



    def mplugImage(self, url):
        pipeline_caption = pipeline(Tasks.image_captioning, model=self.modelName, device=self.deviceType)
        result = pipeline_caption(url)
        str = ''.join(result['caption'].split())
        return str.replace("。", "")
